import os
import json
from pathlib import Path
import shutil

# 输入目录
posenv_ori_dir = '/mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_ori'
posenv_tar_dir = '/mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_tar'

# 输出根目录
style='refcontrolPoseEnv'
output_root = f'/mnt/nas/shengjie/datasets/KontextRefControl_poseenv/{style}'
images_out = Path(output_root) / "images"
images_out.mkdir(parents=True, exist_ok=True)

# 固定 prompt
prompt = "[refcontrolPoseEnv] change pose and env to photo with reference from left side,change person to photo with refence from right side."

# jsonl 输出路径
jsonl_path = Path(output_root) / "train.jsonl"

with open(jsonl_path, "w", encoding="utf-8") as f:
    for filename in sorted(os.listdir(posenv_ori_dir)):
        ori_path = Path(posenv_ori_dir) / filename
        tar_path = Path(posenv_tar_dir) / filename
        if not tar_path.exists():
            continue  # 如果目标图像不存在，跳过

        stem = Path(filename).stem
        src_name = f"{stem}_src{Path(filename).suffix}"
        tar_name = f"{stem}_tar{Path(filename).suffix}"

        # 复制文件到目标目录
        shutil.copy(ori_path, images_out / src_name)
        shutil.copy(tar_path, images_out / tar_name)

        # 写入一条 jsonl
        sample = {
            "src": f"{style}/images/{src_name}",
            "tar": f"{style}/images/{tar_name}",
            "prompt": prompt
        }
        f.write(json.dumps(sample, ensure_ascii=False) + "\n")

print(f"处理完成 ✅ 输出目录: {output_root}")
